from IdentityBasedEncryption import *
from Crypto.Random import get_random_bytes

class FullIdent(IdentityBasedEncryption):
    def __init__(self, lm = 16):
        super(FullIdent, self).__init__()
        if lm not in [16, 32]:
            raise ValueError('only support 128 bits (16 bytes) or 256 bits (32 bytes) message')
        else:
            self.__lm = lm
            self.__hash = md5 if lm == 16 else sha256
            self.__mask = b'\xff' * lm

    def __h1(self, b):
        """ {0, 1} ^ * -> G1 """
        return self.from_hash(b)

    def __h2(self, gt):
        """ {0, 1} ^ * -> G1 """
        bgt = str(gt).encode('utf-8')
        return bytes.fromhex(self.__hash(bgt).hexdigest())

    def __h3(self, b1, b2):
        """ {0, 1} ^ lm X {0, 1} ^ lm -> ZR """
        return self.group.hash(bxor(b1, b2).hex(), ZR)

    def __h4(self, b):
        """ {0, 1} ^ lm -> {0, 1} ^ lm """
        return bxor(b, self.__mask)

    def system_build(self):
        s = self.group.random(ZR)
        return {'msk': s, 'mpk': s * self.g}

    def key_extract(self, user_id, system_params):
        Q_U = self.__h1(user_id)
        return {'usk': system_params['msk'] * Q_U, 'upk': Q_U}

    def encryption(self, m, system_params, user_params):
        if len(m) != self.__lm:
            raise ValueError(f'the bytes of m not satisfied the required length {self.__lm}')
        alpha = get_random_bytes(self.__lm)
        r = self.__h3(alpha, m)
        return r * self.g, bxor(alpha, self.__h2(e(system_params['mpk'], user_params['upk']) ** r)), bxor(m, self.__h4(alpha))

    def decryption(self, ciphertext, user_params):
        V, W, T = ciphertext
        alpha = bxor(W, self.__h2(e(V, user_params['usk'])))
        m = bxor(T, self.__h4(alpha))
        r = self.__h3(alpha, m)
        return m if V == r * self.g else None

if __name__ == '__main__':
    fd = FullIdent()
    fd.test(b'0' * 16, b'xxx')
